import airsim
import math
import numpy as np

objects_dict = {
    "turbine1": "BP_Wind_Turbines_C_1",
    "turbine2": "StaticMeshActor_2",
    "solarpanels": "StaticMeshActor_146",
    "crowd": "StaticMeshActor_6",
    "car": "StaticMeshActor_10",
    "tower1": "SM_Electric_trellis_179",
    "tower2": "SM_Electric_trellis_7",
    "tower3": "SM_Electric_trellis_8",
}


class AirSimWrapper1Uav:
    def __init__(self):
        self.client = airsim.MultirotorClient()
        self.client.confirmConnection()
        self.client.enableApiControl(True)
        self.client.armDisarm(True)

    def takeoff(self):
        self.client.takeoffAsync().join()

    def land(self):
        self.client.landAsync().join()

    def get_drone_position(self):
        pose = self.client.simGetVehiclePose()
        return [pose.position.x_val, pose.position.y_val, pose.position.z_val]

    def fly_to(self, point):
        if point[2] > 0:
            self.client.moveToPositionAsync(point[0], point[1], -point[2], 2).join()
        else:
            self.client.moveToPositionAsync(point[0], point[1], point[2], 2).join()

    def fly_path(self, points):
        airsim_points = []
        for point in points:
            if point[2] > 0:
                airsim_points.append(airsim.Vector3r(point[0], point[1], -point[2]))
            else:
                airsim_points.append(airsim.Vector3r(point[0], point[1], point[2]))
        self.client.moveOnPathAsync(airsim_points, 5, 120, airsim.DrivetrainType.ForwardOnly, airsim.YawMode(False, 0), 20, 1).join()

    def set_yaw(self, yaw):
        self.client.rotateToYawAsync(yaw, 5).join()

    def get_yaw(self):
        orientation_quat = self.client.simGetVehiclePose().orientation
        yaw = airsim.to_eularian_angles(orientation_quat)[2]
        return yaw

    def get_position(self, object_name):
        query_string = objects_dict[object_name] + ".*"
        object_names_ue = []
        while len(object_names_ue) == 0:
            object_names_ue = self.client.simListSceneObjects(query_string)
        pose = self.client.simGetObjectPose(object_names_ue[0])
        return [pose.position.x_val, pose.position.y_val, pose.position.z_val]

    def get_distance(self,  vehicle_name="Drone1"):
       dis =  self.client.getDistanceSensorData(vehicle_name=vehicle_name).distance
       return dis

import time
class AirSimWrapper3Uav:
    def __init__(self):
        self.client = airsim.MultirotorClient()
        self.client.confirmConnection()
        # self.client.enableApiControl(True)
        # self.client.armDisarm(True)
        self. client.enableApiControl(True, "Drone1")
        self.client.enableApiControl(True, "Drone2")
        self.client.enableApiControl(True, "Drone3")
        self.client.armDisarm(True, "Drone1")
        self.client.armDisarm(True, "Drone2")
        self.client.armDisarm(True, "Drone3") 
        # self.get_drone_position1()        
        time.sleep(1)
        self.takeoff() 
        self.arange_flag = True

    def takeoff(self):
        f1 = self.client.takeoffAsync(vehicle_name="Drone1")
        f2 = self.client.takeoffAsync(vehicle_name="Drone2")
        f3 = self.client.takeoffAsync(vehicle_name="Drone3")
        f1.join()
        f2.join()
        f3.join()


    def land(self):
        f1 = self.client.landAsync(vehicle_name="Drone1")
        f2 = self.client.landAsync(vehicle_name="Drone2")
        f3 = self.client.landAsync(vehicle_name="Drone3")
        f1.join()
        f2.join()
        f3.join()

    def get_drone_position(self):
        pose = self.client.simGetVehiclePose(vehicle_name="Drone1")
        return [pose.position.x_val, pose.position.y_val, pose.position.z_val]

    def get_drone_position1(self):
        pose = self.client.simGetVehiclePose(vehicle_name="Drone1")
        pose1 = self.client.simGetVehiclePose(vehicle_name="Drone2")
        pose2 = self.client.simGetVehiclePose(vehicle_name="Drone3")
        print("p1 = ", pose)
        print("p2 = ", pose1)
        print("p3 = ",  pose2)

    def fly_to(self, point):
        if self.arange_flag == True:
            if point[2] > 0:
                f1 = self.client.moveToPositionAsync(point[0], point[1], -point[2], 5, vehicle_name="Drone1")
                f2 = self.client.moveToPositionAsync(point[0], point[1], -point[2], 5, vehicle_name="Drone2")
                f3= self.client.moveToPositionAsync(point[0], point[1], -point[2], 5, vehicle_name="Drone3")
            else:
                f1 = self.client.moveToPositionAsync(point[0], point[1], point[2], 5, vehicle_name="Drone1")
                f2 = self.client.moveToPositionAsync(point[0], point[1] , point[2], 5, vehicle_name="Drone2")
                f3 = self.client.moveToPositionAsync(point[0], point[1], point[2], 5, vehicle_name="Drone3")
            f1.join()
            f2.join()
            f3.join()

        else:        
            if point[2] > 0:
                f1 = self.client.moveToPositionAsync(point[0], point[1], -point[2], 5, vehicle_name="Drone1")
                f2 = self.client.moveToPositionAsync(point[0] + 4, point[1] + 2, -point[2], 5, vehicle_name="Drone2")
                f3= self.client.moveToPositionAsync(point[0] + 8, point[1] - 2, -point[2], 5, vehicle_name="Drone3")
            else:
                f1 = self.client.moveToPositionAsync(point[0], point[1], point[2], 5, vehicle_name="Drone1")
                f2 = self.client.moveToPositionAsync(point[0] + 4, point[1] + 2, point[2], 5, vehicle_name="Drone2")
                f3 = self.client.moveToPositionAsync(point[0] + 8, point[1] -2, point[2], 5, vehicle_name="Drone3")
            f1.join()
            f2.join()
            f3.join()

    def fly_path(self, points):
        airsim_points = []
        for point in points:
            if point[2] > 0:
                airsim_points.append(airsim.Vector3r(point[0], point[1], -point[2]))
            else:
                airsim_points.append(airsim.Vector3r(point[0], point[1], point[2]))
        f1 = self.client.moveOnPathAsync(airsim_points, 5, 120, airsim.DrivetrainType.ForwardOnly, airsim.YawMode(False, 0), 20, 1, vehicle_name="Drone1")
        f2 = self.client.moveOnPathAsync(airsim_points, 5, 120, airsim.DrivetrainType.ForwardOnly, airsim.YawMode(False, 0), 20, 1, vehicle_name="Drone2")
        f3 = self.client.moveOnPathAsync(airsim_points, 5, 120, airsim.DrivetrainType.ForwardOnly, airsim.YawMode(False, 0), 20, 1, vehicle_name="Drone3")
        f1.join()
        f2.join()
        f3.join()

    def set_yaw(self, yaw):
        f1 = self.client.rotateToYawAsync(yaw, 5, vehicle_name='Drone1')
        f2 = self.client.rotateToYawAsync(yaw, 5, vehicle_name='Drone2')
        f3 = self.client.rotateToYawAsync(yaw, 5, vehicle_name='Drone3')
        f1.join()
        f2.join()
        f3.join()


    def get_yaw(self):
        orientation_quat = self.client.simGetVehiclePose(vehicle_name='Drone1').orientation
        yaw = airsim.to_eularian_angles(orientation_quat)[2]
        return yaw

    def get_position(self, object_name):
        query_string = objects_dict[object_name] + ".*"
        object_names_ue = []
        while len(object_names_ue) == 0:
            object_names_ue = self.client.simListSceneObjects(query_string)
        pose = self.client.simGetObjectPose(object_names_ue[0])
        return [pose.position.x_val, pose.position.y_val, pose.position.z_val]


    def arrage_in_column(self):
        if self.arange_flag == True:
            self.arange_flag = False
            pos = self.get_drone_position()
            newpos1 = [pos[0] + 4, pos[1] + 2,  pos[2] + 0.45]
            newpos2 = [pos[0] + 8, pos[1] -  2,  pos[2] + 0.45]
            # print("p1 = ", newpos1)
            # print("p2 = ", newpos2)

            if pos[2] > 0:
                self.client.moveToPositionAsync(newpos2[0], newpos2[1], newpos2[2], 5, vehicle_name="Drone3").join()
                self.client.moveToPositionAsync(newpos1[0], newpos1[1], newpos1[2], 5, vehicle_name="Drone2").join()
            else:
                self.client.moveToPositionAsync(newpos2[0], newpos2[1], newpos2[2], 5, vehicle_name="Drone3").join()
                self.client.moveToPositionAsync(newpos1[0], newpos1[1], newpos1[2], 5, vehicle_name="Drone2").join()
            
            time.sleep(1)    
                
    def arrage_in_row(self):
        if self.arange_flag == False:
            self.arange_flag = True
            pos = self.get_drone_position()
            # print("xxxxxxxx = ", pos)
            newpos1 = [pos[0], pos[1], pos[2]+ 0.45]
            newpos2 = [pos[0], pos[1] , pos[2]+ 0.45]
            # print("p11 = ", newpos1)
            # print("p22 = ", newpos2)            

            if pos[2] > 0:
                f1 = self.client.moveToPositionAsync(newpos1[0], newpos1[1], newpos1[2], 5, vehicle_name="Drone2")
                f2 = self.client.moveToPositionAsync(newpos2[0], newpos2[1], newpos2[2], 5, vehicle_name="Drone3")
                f1.join()
                f2.join()
            else:
                f1 = self.client.moveToPositionAsync(newpos1[0], newpos1[1], newpos1[2], 5, vehicle_name="Drone2")
                f2 = self.client.moveToPositionAsync(newpos2[0], newpos2[1], newpos2[2], 5, vehicle_name="Drone3")
                f1.join()
                f2.join() 
        time.sleep(1)               
             